#!/usr/bin/python3
# -*- coding: utf-8 -*-
import os
import struct

base_path = r'/rootfs/media/kasim/DataSet/Record'
out_path = r'/rootfs/media/kasim/DataSet/Record'
# base_path = r'/rootfs/media/kasim/DataSet/RecordEx'
# out_path = r'/rootfs/media/kasim/DataSet/RecordEx'
video_list_bbox_count_file = r'Face_video_list_bbox_count.txt'
video_list_bbox_file = r'Face_video_list_bbox.dat'
video_list_filter_file = r'video.select.txt'

select_class = 'Face'

select_class_label = {
    'Person': 0,
    'Cat': 1,
    'Dog': 2,
    'BabyCar': 3,
    'Face': 4,
}
select_label = select_class_label[select_class]

THRESHOLDS = [
    0.6,  # Person
    0.6,  # Cat
    0.6,  # Dog
    0.6,  # BabyCar
    0.6,  # Face
]


def main():
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    os.system('chmod a+wr {}'.format(out_path))

    list_file_path = os.path.join(out_path, '{}_video_list_bbox_count.txt'.format(select_class.lower()))
    bbox_path = os.path.join(out_path, '{}_video_list_bbox.dat'.format(select_class.lower()))

    list_file = open(list_file_path, 'w')
    bbox_file = open(bbox_path, 'wb')

    try:
        filter_video_dict = set()
        with open(os.path.join(base_path, video_list_filter_file), 'r') as file:
            for video_info in file.readlines():
                video_name = video_info.split()[0]
                filter_video_dict.add(video_name)
        print('Filter File Count: {}'.format(len(filter_video_dict)))

        video_list = []
        with open(os.path.join(base_path, video_list_bbox_count_file), 'r') as file:
            for video_info in file.readlines():
                video_info = video_info.split(',')
                # video_list.append([video_info[0].strip(), int(video_info[1].strip())])
                video_name = video_info[0].strip()
                video_list.append([video_name, int(video_info[1].strip())])
        print('Read File BBox Count Finish!')

        score_thr = THRESHOLDS[select_label]
        file_count = 0
        total_file_count = len(video_list)
        with open(os.path.join(base_path, video_list_bbox_file), 'rb') as file:
            for video_name, BBox_count in video_list:
                file_count += 1
                if file_count % 1000 == 0:
                    print('File Count: {}/{}'.format(file_count, total_file_count))
                if video_name not in filter_video_dict:
                    bbox_info = file.read(28*BBox_count)
                    if len(bbox_info) < 28*BBox_count:
                        print('File Count: {}/{}'.format(file_count, total_file_count))
                        return
                    continue
                bbox_infos = []
                for j in range(BBox_count):
                    bbox_info = file.read(28)
                    if len(bbox_info) < 28:
                        print('File Count: {}/{}'.format(file_count, total_file_count))
                        return
                    bbox = struct.unpack('6i1f', bbox_info)
                    label = bbox[1]
                    if select_label != label:
                        continue
                    # frame_id = bbox[0]
                    score = bbox[6]
                    if score < score_thr:
                        continue
                    bbox_infos.append(bbox_info)

                if len(bbox_infos) > 0:
                    bbox_count = len(bbox_infos)
                    if bbox_count > 0:
                        file_info_str = '{},{}\n'.format(video_name, bbox_count)
                        list_file.write(file_info_str)
                        for bbox_info in bbox_infos:
                            bbox_file.write(bbox_info)
    except Exception as e:
        print(e)
    print('File Count: {}/{}'.format(file_count, total_file_count))

    list_file.close()
    bbox_file.close()
    os.system('chmod a+wr {}'.format(list_file_path))
    os.system('chmod a+wr {}'.format(bbox_path))

    print('finish!!!!')


if __name__ == '__main__':
    main()
